"""
基于numpy.ndarray 的技术分析指标
输入的数据都为 ndarray
输出是一个或多个 ndarray

"""

import numpy as np
import pandas as pd


# ------------------ 0级：核心工具函数 --------------------------------------------
def RD(N: np.ndarray, D: int = 3):
    """ 四舍五入取 D 位小数"""
    return np.round(N, D)

rd=RD
def RET(S, N=1):  return np.array(S)[-N]  # 返回序列倒数第N个值,默认返回最后一个

ref=RET

def ABS(S):      return np.abs(S)  # 返回N的绝对值


def MAX(S1, S2):  return np.maximum(S1, S2)  # 序列max


def MIN(S1, S2):  return np.minimum(S1, S2)  # 序列min


def MA(S, N):  # 求序列的N日平均值，返回序列
    return pd.Series(S).rolling(N).mean().values


def REF(S, N=1):  # 对序列整体下移动N,返回序列(shift后会产生NAN)
    return pd.Series(S).shift(N).values

def SHIFT(S, N=1):  # 对序列整体下移动N,返回序列(shift后会产生NAN)
    return pd.Series(S).shift(N).values
def DIFF(S, N=1):  # 前一个值减后一个值,前面会产生nan
    return pd.Series(S).diff(N)  # np.diff(S)直接删除nan，会少一行

def TNR(S,N=5):
    """TREND TO NOISE RATIO"""
    D=S.diff(N).abs()
    SD=S.diff(1).abs().sum(N)


def STD(S, N):  # 求序列的N日标准差，返回序列
    return pd.Series(S).rolling(N).std(ddof=0).values


def IF(S_BOOL, S_TRUE, S_FALSE):  # 序列布尔判断 res=S_TRUE if S_BOOL==True  else  S_FALSE
    return np.where(S_BOOL, S_TRUE, S_FALSE)


def SUM(S, N):  # 对序列求N天累计和，返回序列
    return pd.Series(S).rolling(N).sum().values


def HHV(S, N):  # HHV(C, 5)  # 最近5天收盘最高价
    return pd.Series(S).rolling(N).max().values


def LLV(S, N):  # LLV(C, 5)  # 最近5天收盘最低价
    return pd.Series(S).rolling(N).min().values


def EMA(S, N):  # 指数移动平均,为了精度 S>4*N  EMA至少需要120周期
    return pd.Series(S).ewm(span=N, adjust=False).mean().values


def SMA(S, N, M=1):  # 中国式的SMA,至少需要120周期才精确
    K = pd.Series(S).rolling(N).mean()  # 先求出平均值 (下面如果有不用循环的办法，能提高性能，望告知)
    for i in range(N + 1, len(S)):  K[i] = (M * S[i] + (N - M) * K[
        i - 1]) / N  # 因为要取K[i-1]，所以 range(N+1, len(S))
    return K


def AVEDEV(S, N):  # 平均绝对偏差  (序列与其平均值的绝对差的平均值)
    avedev = pd.Series(S).rolling(N).apply(lambda x: (np.abs(x - x.mean())).mean())
    return avedev.values


def SLOPE(S, N, RS=False):  # 返S序列N周期回线性回归斜率 (默认只返回斜率,不返回整个直线序列)
    M = pd.Series(S[-N:]);
    poly = np.polyfit(M.index, M.values, deg=1);
    Y = np.polyval(poly, M.index);
    if RS: return Y[1] - Y[0], Y
    return Y[1] - Y[0]


# ------------------   1级：应用层函数(通过0级核心函数实现） ----------------------------------
def COUNT(S_BOOL, N):  # COUNT(CLOSE>O, N):  最近N天满足S_BOO的天数  True的天数
    return SUM(S_BOOL, N)


def EVERY(S_BOOL, N):  # EVERY(CLOSE>O, 5)   最近N天是否都是True
    R = SUM(S_BOOL, N)
    return IF(R == N, True, False)


def LAST(S_BOOL, A, B):  # 从前A日到前B日一直满足S_BOOL条件
    if A < B: A = B  # 要求A>B    例：LAST(CLOSE>OPEN,5,3)  5天前到3天前是否都收阳线
    return S_BOOL[-A:-B].sum() == (A - B)  # 返回单个布尔值


def EXIST(S_BOOL, N=5):  # EXIST(CLOSE>3010, N=5)  n日内是否存在一天大于3000点
    R = SUM(S_BOOL, N)
    return IF(R > 0, True, False)


def BARSLAST(S_BOOL):  # 上一次条件成立到当前的周期
    M = np.argwhere(S_BOOL);  # BARSLAST(CLOSE/REF(CLOSE)>=1.1) 上一次涨停到今天的天数
    return len(S_BOOL) - int(M[-1]) - 1 if M.size > 0 else -1


def FORCAST(S, N):  # 返S序列N周期回线性回归后的预测值
    K, Y = SLOPE(S, N, RS=True)
    return Y[-1] + K


def CROSS(S1, S2):  # 判断穿越 CROSS(MA(C,5),MA(C,10))
    CROSS_BOOL = IF(S1 > S2, True, False)
    return COUNT(CROSS_BOOL > 0, 2) == 1  # 上穿：昨天0 今天1   下穿：昨天1 今天0


# ------------------   2级：技术指标函数(全部通过0级，1级函数实现） ------------------------------
def MACD(CLOSE, SHORT=12, LONG=26, M=9):  # EMA的关系，S取120日，和雪球小数点2位相同
    DIF = EMA(CLOSE, SHORT) - EMA(CLOSE, LONG);
    DEA = EMA(DIF, M);
    MACD = (DIF - DEA) * 2
    return RD(DIF), RD(DEA), RD(MACD)


def KDJ(CLOSE, HIGH, LOW, N=9, M1=3, M2=3):  # KDJ指标
    RSV = (CLOSE - LLV(LOW, N)) / (HHV(HIGH, N) - LLV(LOW, N)) * 100
    K = EMA(RSV, (M1 * 2 - 1));
    D = EMA(K, (M2 * 2 - 1));
    J = K * 3 - D * 2
    return K, D, J


def RSI(CLOSE, N=24):
    DIF = CLOSE - REF(CLOSE, 1)
    return RD(SMA(MAX(DIF, 0), N) / SMA(ABS(DIF), N) * 100)


def WR(CLOSE, HIGH, LOW, N=10, N1=6):  # W&R 威廉指标
    WR = (HHV(HIGH, N) - CLOSE) / (HHV(HIGH, N) - LLV(LOW, N)) * 100
    WR1 = (HHV(HIGH, N1) - CLOSE) / (HHV(HIGH, N1) - LLV(LOW, N1)) * 100
    return RD(WR), RD(WR1)


def BIAS(CLOSE, L1=6, L2=12, L3=24):  # BIAS乖离率
    BIAS1 = (CLOSE - MA(CLOSE, L1)) / MA(CLOSE, L1) * 100
    BIAS2 = (CLOSE - MA(CLOSE, L2)) / MA(CLOSE, L2) * 100
    BIAS3 = (CLOSE - MA(CLOSE, L3)) / MA(CLOSE, L3) * 100
    return RD(BIAS1), RD(BIAS2), RD(BIAS3)


def BOLL(CLOSE, N=20, P=2):  # BOLL指标，布林带
    MID = MA(CLOSE, N);
    UPPER = MID + STD(CLOSE, N) * P
    LOWER = MID - STD(CLOSE, N) * P
    return RD(UPPER), RD(MID), RD(LOWER)


def PSY(CLOSE, N=12, M=6):
    PSY = COUNT(CLOSE > REF(CLOSE, 1), N) / N * 100
    PSYMA = MA(PSY, M)
    return RD(PSY), RD(PSYMA)


def CCI(CLOSE, HIGH, LOW, N=14):
    TP = (HIGH + LOW + CLOSE) / 3
    return (TP - MA(TP, N)) / (0.015 * AVEDEV(TP, N))


def ATR(CLOSE, HIGH, LOW, N=20):  # 真实波动N日平均值
    TR = MAX(MAX((HIGH - LOW), ABS(REF(CLOSE, 1) - HIGH)), ABS(REF(CLOSE, 1) - LOW))
    return MA(TR, N)


def BBI(CLOSE, M1=3, M2=6, M3=12, M4=20):  # BBI多空指标
    return (MA(CLOSE, M1) + MA(CLOSE, M2) + MA(CLOSE, M3) + MA(CLOSE, M4)) / 4


def DMI(CLOSE, HIGH, LOW, M1=14, M2=6):  # 动向指标：结果和同花顺，通达信完全一致
    TR = SUM(MAX(MAX(HIGH - LOW, ABS(HIGH - REF(CLOSE, 1))), ABS(LOW - REF(CLOSE, 1))), M1)
    HD = HIGH - REF(HIGH, 1);
    LD = REF(LOW, 1) - LOW
    DMP = SUM(IF((HD > 0) & (HD > LD), HD, 0), M1)
    DMM = SUM(IF((LD > 0) & (LD > HD), LD, 0), M1)
    PDI = DMP * 100 / TR;
    MDI = DMM * 100 / TR
    ADX = MA(ABS(MDI - PDI) / (PDI + MDI) * 100, M2)
    ADXR = (ADX + REF(ADX, M2)) / 2
    return PDI, MDI, ADX, ADXR


def TAQ(HIGH, LOW, N):  # 唐安奇通道交易指标，大道至简，能穿越牛熊
    UP = HHV(HIGH, N);
    DOWN = LLV(LOW, N);
    MID = (UP + DOWN) / 2
    return UP, MID, DOWN


def TRIX(CLOSE, M1=12, M2=20):  # 三重指数平滑平均线
    TR = EMA(EMA(EMA(CLOSE, M1), M1), M1)
    TRIX = (TR - REF(TR, 1)) / REF(TR, 1) * 100
    TRMA = MA(TRIX, M2)
    return TRIX, TRMA


def VR(CLOSE, VOL, M1=26):  # VR容量比率
    LC = REF(CLOSE, 1)
    return SUM(IF(CLOSE > LC, VOL, 0), M1) / SUM(IF(CLOSE <= LC, VOL, 0), M1) * 100


def EMV(HIGH, LOW, VOL, N=14, M=9):  # 简易波动指标
    VOLUME = MA(VOL, N) / VOL;
    MID = 100 * (HIGH + LOW - REF(HIGH + LOW, 1)) / (HIGH + LOW)
    EMV = MA(MID * VOLUME * (HIGH - LOW) / MA(HIGH - LOW, N), N);
    MAEMV = MA(EMV, M)
    return EMV, MAEMV


def DPO(CLOSE, M1=20, M2=10, M3=6):  # 区间震荡线
    DPO = CLOSE - REF(MA(CLOSE, M1), M2);
    MADPO = MA(DPO, M3)
    return DPO, MADPO


def BRAR(OPEN, CLOSE, HIGH, LOW, M1=26):  # BRAR-ARBR 情绪指标
    AR = SUM(HIGH - OPEN, M1) / SUM(OPEN - LOW, M1) * 100
    BR = SUM(MAX(0, HIGH - REF(CLOSE, 1)), M1) / SUM(MAX(0, REF(CLOSE, 1) - LOW), M1) * 100
    return AR, BR


def DMA(CLOSE, N1=10, N2=50, M=10):  # 平行线差指标
    DIF = MA(CLOSE, N1) - MA(CLOSE, N2);
    DIFMA = MA(DIF, M)
    return DIF, DIFMA


def MTM(CLOSE, N=12, M=6):  # 动量指标
    MTM = CLOSE - REF(CLOSE, N);
    MTMMA = MA(MTM, M)
    return MTM, MTMMA


def ROC(CLOSE, N=12, M=6):  # 变动率指标
    ROC = 100 * (CLOSE - REF(CLOSE, N)) / REF(CLOSE, N);
    MAROC = MA(ROC, M)
    return ROC, MAROC

    # 望大家能提交更多指标和函数  https://github.com/mpquant/MyTT


class PV4Type:
    """
    连续4峰谷PV4的20种组合: <11, 以谷结束，>10, 以峰结束
    G=GREAT
    E=EQUAL
    L=LITTLE
    P=PEAK
    V=VALLEY
    理论上：
1. 这20种分类是PV4的一种完整划分，记为$$C1=\{ c_i＝i,i=1,\cdots, 20 \}$$ 即
. 任何连续4个峰谷（2峰2谷）属于且只属于这20中的一种；
2. 如果将1、2两种分类合并，11、12合并，形成18种类型，这也是一个完整的划分，记为
$$C2=\{ c_{2,i},i=1,\cdots, 18 \}$$
3. 前10（C2的9）种类型只能来自后10（9）种类型，也只能衍生出后10（9）种类型；
4. 如果将1、2、3、4合并，5、6、7合并，8、9、10合并，11、12、13、14合并，
15、16、17合并，18、19、20合并也将形成一个完整划分，记为
$$C3=\{ c_{3,i},i=1,\cdots, 6 \}$$
5. 将1、2、5、8合并，3、6、9合并，4、7、10合并，11、12、15、18合并，
13、16、19合并，14、17、20合并也将是一个完整划分，记为
$$C4=\{ c_{4,i},i=1,\cdots, 6 \}$$
C3与C4的关系：
$$c_{4,i}将形成c_{3,7-i}，对\forall i\in \{1,\cdots, 6\}$$
如果考虑时间序列，即将连续的$PV=\{PV_t,t=0,\cdots n\}$按C1进行递推分类，形成
$$C1=\{ c_{1,i,t},i\in[1,\cdots, 20], t=4,\cdots, n\}$$
若对$\forall t_0\in [4,\cdots,n-1],c_{1,x,t_0}\in c_{4,i},\exists i \in [1,\cdots,6]$，则：
$$t=t_0+1时,c_{1,y,t}\in c_{3,7-i}$$
成立。

    """
    PGPVGP = 1  # long trend
    PGPVGV = 2
    PGPVEV = 3
    PGPVLV = 4  # short trend
    PEPVGV = 5
    PEPVEV = 6
    PEPVLV = 7  # short trend
    PLPVGV = 8
    PLPVEV = 9
    PLPVLV = 10  # short trend
    VLVPLV = 11
    VLVPLP = 12
    VLVPEP = 13
    VLVPGP = 14  # long trend
    VEVPLP = 15
    VEVPEP = 16
    VEVPGP = 17  # long trend
    VGVPLP = 18
    VGVPEP = 19
    VGVPGP = 20  # long trend

    LONG_COND_PV4 = [1, 2, 3, 5, 6, 8, 9, 14, 17, 20]
    SHORT_COND_PV4 = [11, 12, 13, 15, 16, 18, 19, 4, 7, 10]
    LONG_STOP_COND_ZC = [11, 12, 13, 14, 15, 16, 17, 18, 19, 20]  # 当开多时的 zig_cycle 在此列表中时要设置止损
    SHORT_STOP_COND_ZC = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]  # 当开空时的 zig_cycle 在此列表中时要设置止损
    LONG_STOP_COND_PV = [1]  # 当峰被发现时开多要止损
    SHORT_STOP_COND_PV = [-1]  # 当谷被发现时开空要止损
    C1 = [[1, 2], [3, ], [4, ], [5, ], [6, ], [7, ], [8, ], [9, ], [10, ],
          [11, 12], [13, ], [14, ], [15, ], [16, ], [17, ], [18, ], [19, ], [20, ]]
    C2 = [[1, 2, 3, 4], [5, 6, 7], [8, 9, 10], [11, 12, 13, 14], [15, 16, 17], [18, 19, 20]]
    C3 = [[1, 2, 5, 8], [3, 6, 9], [4, 7, 10], [11, 12, 15, 18], [13, 16, 19], [14, 17, 20]]
    C4 = [[1, 2, 3, 4, 5, 6, 7, 8, 9, 10], [11, 12, 13, 14, 15, 16, 17, 18, 19, 20]]


class PriceClass:
    """价格走势分类定义"""
    '''盘整：当时间长度大于给定的时间窗口T（序列长度）时，价格的波动区间小于给定的幅度R（百分比）'''
    PZ = 0
    '''升开始：向上脱离盘整区间，或下降结束转上升，开始创新高，在给定的时间窗口T内，价格波动幅度不小于给定的R；最好的做多区间，次优的平空区间'''
    XB = 1
    '''升持续：反复创新高；持多仓'''
    XL = 4
    '''升结束：在给定的时间窗口T内不再创新高，可转为盘整或降开始；最好的平多区间，次优的开空区间'''
    XE = 8
    '''降开始：向下脱离盘整区间，或上升结束转下降，开始创新低；最优的开空区间，次优的平多区间'''
    JB = -1
    '''降持续：反复创新低；持空仓'''
    JL = -4
    '''降结束：在给定的时间窗口T内，不再创新低，盘整开始，或转上升；最优的平空区间，次优的开多区间'''
    JE = -8


def tick_zig(last_price: np.array, r=20, start_posi=0, old_zig=None):
    """
    注意：这个算法的最后一个峰/谷是未确定的，在实际应用时要另行处理
    :param last_price: np.array 一个一维数组
    :param r: 要求的最小绝对差值
    :param start_posi:
    :param old_zig: 已经计算好的zig，只有最后一个数据没有参与计算，＝＝None时表示没有计算
    :return: nd.array([[flag,pv,position]])
    """
    if old_zig is not None:
        current_zig = old_zig[-1]
        flag, value, posi = current_zig
        part_data = last_price[posi:]
        new_zig = tick_zig(part_data, r)
        if new_zig is None:
            return old_zig
        if len(new_zig) >= 2:
            tmp = list(old_zig)
            tmp[-1] = [flag, value, posi]
            tmp.extend(new_zig[1:])
            return np.array(tmp)
        else:
            old_zig[-1] = new_zig[0]
            return old_zig

    p, pp = last_price.max(), last_price.argmax()  # peak,peak position
    v, vp = last_price.min(), last_price.argmin()  # valley,valley position
    if p - v > r:
        if pp > vp:  # 先谷后峰
            pv = [[-1, v, vp + start_posi]]
            # 谷前部分
            if vp > 0:
                new_pv = tick_zig(last_price[0:vp + 1], r, start_posi)
                if new_pv is not None:  # 最后一个是已经在pv中的谷，不要再次加入
                    for i in range(len(new_pv) - 1):
                        pv.insert(i, new_pv[i])
            # 峰谷间部分
            if pp - vp > 2:
                v_o_p = 0
                # 寻找谷的发现点
                for i in range(vp + 1, pp):
                    if last_price[i] - v > r:
                        v_o_p = i
                        break
                if v_o_p > 0 and pp - v_o_p > 1:
                    new_pv = tick_zig(last_price[v_o_p:pp + 1], r, start_posi=v_o_p + start_posi)
                    if new_pv is not None and len(new_pv) > 2:
                        if len(new_pv) % 2 == 0:
                            pv.extend(new_pv[1:-1])
                        else:
                            pv.extend(new_pv[:-1])

            pv.append([1, p, pp + start_posi])
            # 峰后部分
            if pp < len(last_price) - 1:
                new_pv = tick_zig(last_price[pp:], r, start_posi=pp + start_posi)
                if new_pv is not None:
                    pv.extend(new_pv[1:])

        else:  # 先峰后谷
            pv = [[1, p, pp + start_posi]]
            # 峰前部分
            if pp > 0:
                new_pv = tick_zig(last_price[0:pp + 1], r, start_posi=start_posi)
                if new_pv is not None:
                    for i in range(len(new_pv) - 1):
                        pv.insert(i, new_pv[i])
            # 峰谷间部分
            if vp - pp > 1:
                p_o_p = 0
                # 寻找峰的发现点
                for i in range(pp + 1, vp):
                    if p - last_price[i] > r:
                        p_o_p = i
                        break
                if p_o_p > 0 and vp - p_o_p > 1:
                    new_pv = tick_zig(last_price[p_o_p:vp + 1], r, start_posi=p_o_p + start_posi)
                    if new_pv is not None and len(new_pv) > 2:
                        if len(new_pv) % 2 == 0:
                            pv.extend(new_pv[1:-1])
                        else:
                            pv.extend(new_pv[:-1])
            # 加谷
            pv.append([-1, v, vp + start_posi])
            # 谷后部分
            if vp < len(last_price) - 1:
                new_pv = tick_zig(last_price[vp:], r, start_posi=vp + start_posi)
                if new_pv is not None:
                    pv.extend(new_pv[1:])

        return np.array(pv)
    else:
        return None


def zig(h: np.array, l: np.array, r=20, start_posi=0, old_zig=None):
    """
    注意：这个算法的最后峰谷是未确定的，在实际应用时要另行处理
    :param h: np.array 一个一维数组,在同一位置，h>=l
    :param l: np.array 一个一维数组,在同一位置，h>=l
    :param r: 要求的最小绝对差值
    :param start_posi:
    :param old_zig: 已经计算好的zig，只有最后一个数据没有参与计算，＝＝None时表示没有计算
    :return: nd.array([[flag,pv,posi]])
    """
    if old_zig is not None:
        current_zig = old_zig[-1]
        flag, value, posi = current_zig
        part_data = h[posi:], l[posi:]
        new_zig = zig(part_data, r)
        if new_zig is None:
            return old_zig
        if len(new_zig) >= 2:
            tmp = list(old_zig)
            tmp[-1] = [flag, value, posi]
            tmp.extend(new_zig[1:])
            return np.array(tmp)
        else:
            old_zig[-1] = new_zig[0]
            return old_zig

    p, pp = h.max(), h.argmax()
    v, vp = l.min(), l.argmin()
    if p - v > r:
        if pp > vp:  # 先谷后峰
            pv = [[-1, v, vp + start_posi]]
            # 谷前部分
            if vp > 0:
                new_pv = zig(h[0:vp + 1], l[0:vp + 1], r, start_posi)
                if new_pv is not None:  # 最后一个是已经在pv中的谷，不要再次加入
                    for i in range(len(new_pv) - 1):
                        pv.insert(i, new_pv[i])
            # 峰谷间部分
            if pp - vp > 2:
                v_o_p = 0
                # 寻找谷的发现点
                for i in range(vp + 1, pp):
                    if h[i] - v > r:
                        v_o_p = i
                        break
                if v_o_p > 0 and pp - v_o_p > 1:
                    new_pv = zig(h[v_o_p:pp + 1], l[v_o_p:pp + 1], r, start_posi=v_o_p + start_posi)
                    if new_pv is not None and len(new_pv) > 2:
                        if len(new_pv) % 2 == 0:
                            pv.extend(new_pv[1:-1])
                        else:
                            pv.extend(new_pv[:-1])

            pv.append([1, p, pp + start_posi])
            # 谷后部分
            if pp < len(h) - 1:
                new_pv = zig(h[pp:], l[pp:], r, start_posi=pp + start_posi)
                if new_pv is not None:
                    pv.extend(new_pv[1:])

        else:  # 先峰后谷
            pv = [[1, p, pp + start_posi]]
            # 峰前部分
            if pp > 0:
                new_pv = zig(h[0:pp + 1], l[0:pp + 1], r, start_posi=start_posi)
                if new_pv is not None:
                    for i in range(len(new_pv) - 1):
                        pv.insert(i, new_pv[i])
            # 峰谷间部分
            if vp - pp > 1:
                p_o_p = 0
                # 寻找峰的发现点
                for i in range(pp + 1, vp):
                    if p - l[i] > r:
                        p_o_p = i
                        break
                if p_o_p > 0 and vp - p_o_p > 1:
                    new_pv = zig(h[p_o_p:vp + 1], l[p_o_p:vp + 1], r, start_posi=p_o_p + start_posi)
                    if new_pv is not None and len(new_pv) > 2:
                        if len(new_pv) % 2 == 0:
                            pv.extend(new_pv[1:-1])
                        else:
                            pv.extend(new_pv[:-1])
            # 加谷
            pv.append([-1, v, vp + start_posi])
            # 峰后部分
            if vp < len(l) - 1:
                new_pv = zig(h[vp:], l[vp:], r, start_posi=vp + start_posi)
                if new_pv is not None:
                    pv.extend(new_pv[1:])

        return np.array(pv)
    else:
        return None


def clc_zg(p0, v0, h0, l0, flag, r):
    """
    根据当前的峰、谷值，及给定的高低价、和峰谷标志，计算是否产生新的峰谷并返回
    :param p0: 当前的峰值，或者是最高价（flag==2)
    :param v0:当前的谷值，或者是最低价（flag==-2)
    :param h0:新的高价
    :param l0:新的低价
    :param flag: 峰谷标志 1＝＝峰，－1＝＝谷，0＝＝第一个数据，2＝＝高先，－2＝＝低先
    :param r: 确定峰谷要求的绝对距离
    :return: p,v,new_flag,pv,若 pv==0 还没有找到满足条件的峰或谷
    """
    pv = 0
    if flag == 1:  # 当前的峰已经确认
        if l0 < v0:
            return h0, l0, flag, pv
        elif h0 >= p0:
            p0 = h0
            if p0 - v0 >= r:
                return p0, l0, 0 - flag, v0
            else:
                return p0, v0, flag, pv
        else:
            return p0, v0, flag, pv
    elif flag == -1:  # 当前的谷已经确认
        if h0 > p0:
            return h0, l0, flag, pv
        elif l0 <= v0:
            v0 = l0
            if p0 - v0 >= r:
                return h0, v0, 0 - flag, p0
            else:
                return p0, v0, flag, pv
        else:
            return p0, v0, flag, pv
    elif flag == 0:  # 开始峰谷未产生
        v0 = l0
        p0 = h0
        flag = 2  # 也可以高为－2
        return p0, v0, flag, pv
    elif flag == 2:
        if h0 > p0:
            p0 = h0
            if p0 - v0 >= r:
                flag = -1
                return p0, l0, flag, v0
            else:
                flag = -2
                return p0, v0, flag, pv
        elif l0 < v0:
            v0 = l0
            if p0 - v0 >= r:
                flag = 1
                return h0, v0, flag, p0
            else:
                return p0, v0, flag, pv
        else:
            return p0, v0, flag, pv
    elif flag == -2:
        if h0 > p0:
            p0 = h0
            if p0 - v0 >= r:
                flag = -1
                return p0, l0, flag, v0
            else:
                return p0, v0, flag, pv
        elif l0 < v0:
            v0 = l0
            if p0 - v0 >= r:
                flag = 1
                return h0, v0, flag, p0
            else:
                flag = 2
                return p0, v0, flag, pv
        else:
            return p0, v0, flag, pv
    else:
        return p0, v0, flag, pv


def zig_p(h: np.array, l: np.array, r=20, start_posi=0, old_zig=None, center=0):
    """
    注意：这个算法的最后个 峰谷是未确定的，在实际应用时要另行处理
    :param h: np.array 一个一维数组,在同一位置，h>=l
    :param l: np.array 一个一维数组,在同一位置，h>=l
    :param r: 要求的最小相对千分比
    :param start_posi:
    :param old_zig: 已经计算好的zig，只有最后一个数据没有参与计算，＝＝None时表示没有计算
    :param center: 计算 相对比例的分母的方式：0＝(p+v)/2,1=先出现的
    :return: nd.array([[flag,pv,posi]])
    """
    if old_zig is not None:
        current_zig = old_zig[-1]
        flag, value, posi = current_zig
        part_data = h[posi:], l[posi:]
        new_zig = zig(part_data, r)
        if new_zig is None:
            return old_zig
        if len(new_zig) >= 2:
            tmp = list(old_zig)
            tmp[-1] = [flag, value, posi]
            tmp.extend(new_zig[1:])
            return np.array(tmp)
        else:
            old_zig[-1] = new_zig[0]
            return old_zig

    p, pp = h.max(), h.argmax()
    v, vp = l.min(), l.argmin()
    x = r / 1000.  # 临时变量
    if center == 0:
        x *= (p + v) / 2.
    elif pp > vp:
        x *= v
    else:
        x *= p
    if p - v > x:
        if pp > vp:  # 先谷后峰
            pv = [[-1, v, vp + start_posi]]
            # 谷前部分
            if vp > 0:
                new_pv = zig(h[0:vp + 1], l[0:vp + 1], r, start_posi)
                if new_pv is not None:  # 最后一个是已经在pv中的谷，不要再次加入
                    for i in range(len(new_pv) - 1):
                        pv.insert(i, new_pv[i])
            # 峰谷间部分
            if pp - vp > 2:
                v_o_p = 0
                # 寻找谷的发现点
                for i in range(vp + 1, pp):
                    if h[i] - v > p:
                        v_o_p = i
                        break
                if v_o_p > 0 and pp - v_o_p > 1:
                    new_pv = zig(h[v_o_p:pp + 1], l[v_o_p:pp + 1], r, start_posi=v_o_p + start_posi)
                    if new_pv is not None and len(new_pv) > 2:
                        if len(new_pv) % 2 == 0:
                            pv.extend(new_pv[1:-1])
                        else:
                            pv.extend(new_pv[:-1])

            pv.append([1, p, pp + start_posi])
            # 谷后部分
            if pp < len(h) - 1:
                new_pv = zig(h[pp:], l[pp:], r, start_posi=pp + start_posi)
                if new_pv is not None:
                    pv.extend(new_pv[1:])

        else:  # 先峰后谷
            pv = [[1, p, pp + start_posi]]
            # 峰前部分
            if pp > 0:
                new_pv = zig(h[0:pp + 1], l[0:pp + 1], r, start_posi=start_posi)
                if new_pv is not None:
                    for i in range(len(new_pv) - 1):
                        pv.insert(i, new_pv[i])
            # 峰谷间部分
            if vp - pp > 1:
                p_o_p = 0
                # 寻找峰的发现点
                for i in range(pp + 1, vp):
                    if p - l[i] > p:
                        p_o_p = i
                        break
                if p_o_p > 0 and vp - p_o_p > 1:
                    new_pv = zig(h[p_o_p:vp + 1], l[p_o_p:vp + 1], r, start_posi=p_o_p + start_posi)
                    if new_pv is not None and len(new_pv) > 2:
                        if len(new_pv) % 2 == 0:
                            pv.extend(new_pv[1:-1])
                        else:
                            pv.extend(new_pv[:-1])
            # 加谷
            pv.append([-1, v, vp + start_posi])
            # 峰后部分
            if vp < len(l) - 1:
                new_pv = zig(h[vp:], l[vp:], r, start_posi=vp + start_posi)
                if new_pv is not None:
                    pv.extend(new_pv[1:])

        return np.array(pv)
    else:
        return None


def clc_zg_p(p0: float, v0: float, h0: float, l0: float, flag, p: float, center=0):
    """
    根据当前的峰、谷值，及给定的高低价、和峰谷标志，计算是否产生新的峰谷并返回
    :param p0: 当前的峰值，或者是最高价（flag==2)
    :param v0:当前的谷值，或者是最低价（flag==-2)
    :param h0:新的高价
    :param l0:新的低价
    :param flag: 峰谷标志 1＝＝峰，－1＝＝谷，0＝＝第一个数据，2＝＝高先，－2＝＝低先
    :param p: 确定峰谷要求的最小相对千分比
    :return: p,v,new_flag,pv,若 pv==0 还没有找到满足条件的峰或谷
    """

    def clc(p0: float, v0: float, h0: float, l0: float, flag, p: float):
        pv = 0
        if flag == 1:  # 当前的峰已经确认
            if l0 < v0:
                return h0, l0, flag, pv
            elif h0 >= p0:
                p0 = h0
                if p0 - v0 >= p:
                    return p0, l0, 0 - flag, v0
                else:
                    return p0, v0, flag, pv
            else:
                return p0, v0, flag, pv
        elif flag == -1:  # 当前的谷已经确认
            if h0 > p0:
                return h0, l0, flag, pv
            elif l0 <= v0:
                v0 = l0
                if p0 - v0 >= p:
                    return h0, v0, 0 - flag, p0
                else:
                    return p0, v0, flag, pv
            else:
                return p0, v0, flag, pv
        elif flag == 0:  # 开始峰谷未产生
            v0 = l0
            p0 = h0
            flag = 2  # 也可以高为－2
            return p0, v0, flag, pv
        elif flag == 2:
            if h0 > p0:
                p0 = h0
                if p0 - v0 >= p:
                    flag = -1
                    return p0, l0, flag, v0
                else:
                    flag = -2
                    return p0, v0, flag, pv
            elif l0 < v0:
                v0 = l0
                if p0 - v0 >= p:
                    flag = 1
                    return h0, v0, flag, p0
                else:
                    return p0, v0, flag, pv
            else:
                return p0, v0, flag, pv
        elif flag == -2:
            if h0 > p0:
                p0 = h0
                if p0 - v0 >= p:
                    flag = -1
                    return p0, l0, flag, v0
                else:
                    return p0, v0, flag, pv
            elif l0 < v0:
                v0 = l0
                if p0 - v0 >= p:
                    flag = 1
                    return h0, v0, flag, p0
                else:
                    flag = 2
                    return p0, v0, flag, pv
            else:
                return p0, v0, flag, pv
        else:
            return p0, v0, flag, pv

    if center == 0:
        x = (max(p0, h0) + min(v0, l0)) * 2 * p / 1000
        return clc(p0, v0, h0, l0, flag, x)
    elif flag == 1:
        x = v0 * p / 1000
        return clc(p0, v0, h0, l0, flag, x)
    elif flag == -1:
        x = p0 * p / 1000
        return clc(p0, v0, h0, l0, flag, x)
    elif flag == 2:
        x = p0 * p / 1000
        return clc(p0, v0, h0, l0, flag, x)
    else:
        x = v0 * p / 1000
        return clc(p0, v0, h0, l0, flag, x)


def clc_zc(pv0, pv1, pv2, pv3):
    """
    根据给定的两峰两谷的值及其顺序，确定这连续的4个峰谷的类型
    :param pv0: 第一个峰谷
    :param pv1: 第二个峰谷
    :param pv2: 第三个峰谷
    :param pv3: 第四个峰谷
    :return: KLine4Type 中的某一个值
    """
    if pv0 > pv1:  # 峰先
        p0, v0, p1, v1 = pv0, pv1, pv2, pv3

        if p1 > p0:
            if v1 > v0:
                c = PV4Type.PGPVGP if v1 > p0 else PV4Type.PGPVGV
            elif v1 == v0:
                c = PV4Type.PGPVEV
            else:
                c = PV4Type.PGPVLV
        elif p1 == p0:
            if v1 > v0:
                c = PV4Type.PEPVGV
            elif v1 == v0:
                c = PV4Type.PEPVEV
            else:
                c = PV4Type.PEPVLV
        else:
            if v1 > v0:
                c = PV4Type.PLPVGV
            elif v1 == v0:
                c = PV4Type.PLPVEV
            else:
                c = PV4Type.PLPVLV
    else:  # 谷先
        v0, p0, v1, p1 = pv0, pv1, pv2, pv3
        if p1 > p0:
            if v1 > v0:
                c = PV4Type.VGVPGP
            elif v1 == v0:
                c = PV4Type.VEVPGP
            else:
                c = PV4Type.VLVPGP
        elif p1 == p0:
            if v1 > v0:
                c = PV4Type.VGVPEP
            elif v1 == v0:
                c = PV4Type.VEVPEP
            else:
                c = PV4Type.VLVPEP
        else:
            if v1 > v0:
                c = PV4Type.VGVPLP
            elif v1 == v0:
                c = PV4Type.VEVPLP
            else:
                c = PV4Type.VLVPLP if v0 > p1 else PV4Type.VLVPLV
    return c


def zig_cycle(zg, start_pos=0, old_zc=None):
    """定义连续4个峰谷（2峰2谷）[v0,p0,v1,p1] [p0,v0,p1,v1].的关系zc
    :param zg: zig 的计算结果，长度最小为4，
    :param start_pos: 开始计算的位置，0：从头开始，－1：只计算最后一个
    :param old_zc: 已经计算好的，新计算的结果加到后面并返回.
    :return: [PV4Type]

 """

    pv_is = zg[:, 0]
    # pv = zg[:, 1]
    dl = len(pv_is)
    if dl < 4:
        return
    i = start_pos
    zc = [0, 0, 0]
    if i == 0:  # 从第一个开始
        pv0 = zg[i, 1]
        i += 1
        pv1 = zg[i, 1]
        i += 1
        pv2 = zg[i, 1]
        i += 1
        pv3 = zg[i, 1]

        while i < dl:
            c = clc_zc(pv0, pv1, pv2, pv3)
            zc.append(c)
            i += 1

            if i < dl:
                pv0, pv1, pv2, pv3 = pv1, pv2, pv3, zg[i, 1]
            else:
                break

    elif i == -1:  # 最后一个是峰
        pv0, pv1, pv2, pv3 = zg[i - 3:, 1]
        c = clc_zc(pv0, pv1, pv2, pv3)
        zc = list(old_zc)
        zc.append(c)

    zc = np.array(zc)

    return zc  # np.array([zc, zc_pre]).transpose()


def max_drawback(h, l, absolute=True):
    """计算最大回撤（值－绝对值或相对值，最大值位置，最小值位置），包括做多和做空
    :param h: 最高价向量
    :param l: 最低价向量
    :param absolute: ==Ture 时，用绝对值， 否则取相对值
    :return : 做多时的最大回撤，做空时的最大回撤[最大回撤值，最高值位置，最低值位置]

    """
    mn = np.minimum.accumulate(l)
    mx = np.maximum.accumulate(h)
    # 计算做多时的最大回撤（p_mx_d, p_mx_p, p_mn_p）
    if absolute:
        d = mx - l
    else:
        d = (mx - l) / mx
    i = np.argmax(d)
    p_mx_d = d[i]
    p_mn_p = i
    j = np.argmax(h[:i])
    p_mx_p = j

    # 计算做空时的最大回撤（n_mx_d, n_mx_p, n_mn_p）
    if absolute:
        d = h - mn
    else:
        d = (h - mn) / mn
    i = np.argmax(d)
    n_mx_d = d[i]
    n_mx_p = i
    j = np.argmin(l[:i])
    n_mn_p = j

    return p_mx_d, p_mx_p, p_mn_p, n_mx_d, n_mx_p, n_mn_p


def max_drawback_long(h, l, absolute=True):
    """计算做多最大回撤（值－绝对值或相对值，最大值位置，最小值位置）
    :param h: 最高价向量
    :param l: 最低价向量
    :param absolute: ==Ture 时，用绝对值， 否则取相对值
    :return : 做多时的最大回撤，做空时的最大回撤

    """
    mx = np.maximum.accumulate(h)
    # 计算做多时的最大回撤（p_mx_d, p_mx_p, p_mn_p）
    if absolute:
        d = mx - l
    else:
        d = (mx - l) / mx
    i = np.argmax(d)
    p_mx_d = d[i]
    p_mn_p = i
    j = np.argmax(h[:i])
    p_mx_p = j

    return p_mx_d, p_mx_p, p_mn_p


def max_drawback_short(h, l, absolute=True):
    """计算做空最大回撤（值－绝对值或相对值，最大值位置，最小值位置）
    :param h: 最高价向量
    :param l: 最低价向量
    :param absolute: ==Ture 时，用绝对值， 否则取相对值
    :return : 做多时的最大回撤，做空时的最大回撤

    """
    mn = np.minimum.accumulate(l)

    # 计算做空时的最大回撤（n_mx_d, n_mx_p, n_mn_p）
    if absolute:
        d = h - mn
    else:
        d = (h - mn) / mn
    i = np.argmax(d)
    n_mx_d = d[i]
    n_mx_p = i
    j = np.argmin(l[:i])
    n_mn_p = j

    return n_mx_d, n_mx_p, n_mn_p


def hp(y, lamb=10):
    """HP滤波"""

    def D_matrix(N):
        D = np.zeros((N - 1, N))
        D[:, 1:] = np.eye(N - 1)
        D[:, :-1] -= np.eye(N - 1)
        """D=
        [[-1.  1.  0. ...  0.  0.  0.]
         [ 0. -1.  1. ...  0.  0.  0.]
         [ 0.  0. -1. ...  0.  0.  0.]
         ...
         [ 0.  0.  0. ...  1.  0.  0.]
         [ 0.  0.  0. ... -1.  1.  0.]
         [ 0.  0.  0. ...  0. -1.  1.]]
        """
        return D

    N = len(y)
    D1 = D_matrix(N)
    D2 = D_matrix(N - 1)
    D = D2 @ D1
    g = np.linalg.inv((np.eye(N) + lamb * D.T @ D)) @ y
    return g


def brick_k(op, hp, lp, cp, brick_height=0, permill=1, bars=None):
    """
    砖K线
    :param op, hp, lp, cp:是开高低收价格向量，或者是 tick.last_price
    :param brick_height: 计算砖的绝对高度
    :param permill: 计算砖的高度的千分比
    :return:
    """
    b = [[hp[0], lp[0], 0]]  # [b_h,b_l,b_type], b_type:0==None, 1==up, -1==down,
    cnt = len(bars) if bars is not None else len(hp)
    if bars is not None:
        op, hp, lp, cp = bars[:]
    if brick_height == 0:
        x = int((np.log(hp[0] / lp[0]) / np.log(1. - permill / 1000.)))
        if op[0] > cp[0]:
            if x > 0:
                b = [[hp[0], hp[0] * (1. - permill / 1000.), -1]]
                for i in range(1, x):
                    b.append([hp[0] * (1 - permill / 1000.) ^ i, hp[0] * (1 - permill / 1000.) ^ (i + 1), -1])
        else:
            if x > 0:
                b = [[lp[0] * (1. + permill / 1000.), lp[0], 1]]
                for i in range(1, x):
                    b.append([lp[0] * (1 + permill / 1000.) ^ (i + 1), lp[0] * (1 + permill / 1000.) ^ i, 1])
        for i in range(1, cnt):
            h = b[-1][0]
            l = b[-1][1]
            f = b[-1][2]
            x = int(np.log(hp[i] / h) / np.log(1. + permill / 1000.))
            if x > 0:
                b.append([[h * (1 + permill / 1000.), h, 1]])
                for i in range(1, x):
                    b.append([h * (1 + permill / 1000.) ^ (1 + i), h * (1 + permill / 1000.) ^ i, 1])
                continue
            x = int(np.log(l / lp[i]) / np.log(1. - permill / 1000.))
            if x > 0:
                b.append([[l, l * (1 - permill / 1000.), -1]])
                for i in range(1, x):
                    b.append([l * (1 - permill / 1000.) ^ i, l * (1 - permill / 1000.) ^ (i + 1), -1])
                continue
        return b
    else:
        x = int((hp[0] - lp[0]) / brick_height)
        if op[0] > cp[0]:
            if x > 0:
                b = [[hp[0], hp[0] - brick_height, -1]]
                for i in range(1, x):
                    b.append([hp[0] - brick_height * i, hp[0] - brick_height * (i + 1), -1])
        else:
            if x > 0:
                b = [[lp[0] + brick_height, lp[0], 1]]
                for i in range(1, x):
                    b.append([lp[0] + brick_height * (i + 1), lp[0] + brick_height * i, 1])

        for i in range(1, cnt):
            h = b[-1][0]
            l = b[-1][1]
            f = b[-1][2]
            x = int((hp[i] - h) / brick_height)
            if x > 0:
                b.append([[h, h + brick_height, 1]])
                for i in range(1, x):
                    b.append([h + brick_height * i, h + brick_height * (i + 1), 1])
                continue
            x = int((l - lp[i]) / brick_height)
            if x > 0:
                b.append([[l, l - brick_height, -1]])
                for i in range(1, x):
                    b.append([l - brick_height * i, l - brick_height * (i + 1), -1])
                continue
        return b


def zig_pr(h: np.array, l: np.array, pr=0.005, start_posi=0, old_zig=None):
    """
    注意：这个算法的最后峰谷是未确定的，在实际应用时要另行处理
    :param h: np.array 一个一维数组,在同一位置，h>=l
    :param l: np.array 一个一维数组,在同一位置，h>=l
    :param pr: 要求的最小差值的百分比
    :param start_posi:
    :param old_zig: 已经计算好的zig，只有最后一个数据没有参与计算，＝＝None时表示没有计算
    :return: nd.array([[flag,pv,posi]])
    """
    if old_zig is not None:
        current_zig = old_zig[-1]
        flag, value, posi = current_zig
        part_data = h[posi:], l[posi:]
        new_zig = zig(part_data, pr)
        if new_zig is None:
            return old_zig
        if len(new_zig) >= 2:
            tmp = list(old_zig)
            tmp[-1] = [flag, value, posi]
            tmp.extend(new_zig[1:])
            return np.array(tmp)
        else:
            old_zig[-1] = new_zig[0]
            return old_zig

    p, pp = h.max(), h.argmax()
    v, vp = l.min(), l.argmin()
    if p - v > pr * (p + v) / 2:
        if pp > vp:  # 先谷后峰
            pv = [[-1, v, vp + start_posi]]
            # 谷前部分
            if vp > 0:
                new_pv = zig(h[0:vp + 1], l[0:vp + 1], pr, start_posi)
                if new_pv is not None:  # 最后一个是已经在pv中的谷，不要再次加入
                    for i in range(len(new_pv) - 1):
                        pv.insert(i, new_pv[i])
            # 峰谷间部分
            if pp - vp > 2:
                v_o_p = 0
                # 寻找谷的发现点
                for i in range(vp + 1, pp):
                    if h[i] - v > pr * (p + v) / 2:
                        v_o_p = i
                        break
                if v_o_p > 0 and pp - v_o_p > 1:
                    new_pv = zig(h[v_o_p:pp + 1], l[v_o_p:pp + 1], pr, start_posi=v_o_p + start_posi)
                    if new_pv is not None and len(new_pv) > 2:
                        if len(new_pv) % 2 == 0:
                            pv.extend(new_pv[1:-1])
                        else:
                            pv.extend(new_pv[:-1])

            pv.append([1, p, pp + start_posi])
            # 谷后部分
            if pp < len(h) - 1:
                new_pv = zig(h[pp:], l[pp:], pr, start_posi=pp + start_posi)
                if new_pv is not None:
                    pv.extend(new_pv[1:])

        else:  # 先峰后谷
            pv = [[1, p, pp + start_posi]]
            # 峰前部分
            if pp > 0:
                new_pv = zig(h[0:pp + 1], l[0:pp + 1], pr, start_posi=start_posi)
                if new_pv is not None:
                    for i in range(len(new_pv) - 1):
                        pv.insert(i, new_pv[i])
            # 峰谷间部分
            if vp - pp > 1:
                p_o_p = 0
                # 寻找峰的发现点
                for i in range(pp + 1, vp):
                    if p - l[i] > pr * (p + v) / 2:
                        p_o_p = i
                        break
                if p_o_p > 0 and vp - p_o_p > 1:
                    new_pv = zig(h[p_o_p:vp + 1], l[p_o_p:vp + 1], pr, start_posi=p_o_p + start_posi)
                    if new_pv is not None and len(new_pv) > 2:
                        if len(new_pv) % 2 == 0:
                            pv.extend(new_pv[1:-1])
                        else:
                            pv.extend(new_pv[:-1])
            # 加谷
            pv.append([-1, v, vp + start_posi])
            # 峰后部分
            if vp < len(l) - 1:
                new_pv = zig(h[vp:], l[vp:], pr, start_posi=vp + start_posi)
                if new_pv is not None:
                    pv.extend(new_pv[1:])

        return np.array(pv)
    else:
        return None


def zig_windows(high: pd.Series, low: pd.Series, rng_pct=0.1, min_windows: int = 10):
    """
    基于窗口长度的:zig
    :param high: 最高价序列
    :param low: 最低价序列
    :param rng_pct: 要求的最小波动幅度百分比
    :param min_windows:　最小窗口长度
    :return:
    """
    h = high.rolling(min_windows).max()
    l = low.rolling(min_windows).min()
    # 去掉连续相同的，保留第一个
    hh = h[h.diff(1) != 0]
    ll = l[l.diff(1) != 0]
    # 相邻不同高间的距离
    dh = hh.diff(1)
    dl = ll.diff(1)
    # 去掉连续增加或减小的
    sdh = pd.Series([0] * len(dh), index=dh.index)
    sdh[dh > 0] = 1
    sdh[dh < 0] = -1
    sdh0 = sdh[sdh.diff(1) != 0]  # 保留第一个
    sdh_1 = sdh[sdh[::-1].diff(1) != 0]  # 保留最后一个
    sdl = pd.Series([0] * len(dl), index=dl.index)
    sdl[dl > 0] = 1
    sdl[dl < 0] = -1
    sdl0 = sdl[sdl.diff(1) == 0]  # 保留第一个
    sdl_1 = sdl[sdl[::-1].diff(1) == 0]  # 保留最后一个
    return sdh0, sdl0, sdh_1, sdl_1


def ama(price: pd.Series, m: int):
    """计算正常的ma(price,m),从ma的转折点开始重新计算ma"""
    ma = price.rolling(window=m, min_periods=1).mean()
    # ma的差分
    madiff = ma.diff(1)
    # 差分后的符号
    smad = madiff.sgn()


if __name__ == '__main__':
    # a=pd.read_csv('tmp.csv').to_numpy()[:,1]
    r = 10
    a = np.array([90, 80, 70, 60, 1000, 20, 40, 80, 0, 100, 110, 105, 120, 5, 6, 4, 7])
    # a = [90, 80, 70, 60, 1000, 20, 40, 80, 0, 100, 110, 105, 120, 5, 6, 4, 7]
    print(type(a), len(a))
    print(max_drawback(a, a))
    tz = tick_zig(a, 60)
    print(tz)
    b = pd.Series(a)
    zw = zig_windows(b, b, 3)
    print(zw)
